from querier.esquerier import ElasticSearchQuerier


class WeiboDataOverviewQuerier(ElasticSearchQuerier):
    def __init__(self, es, index, doc_type):
        super(WeiboDataOverviewQuerier, self).__init__(es, index, doc_type)

    def _build_query(self, args):
        user_id = args.get('user_id', None)
        from_date = args.get('from', None)
        to_date = args.get('to', None)
        filters = args.get('filters', {})
        filters = filters if filters else {}

        if user_id is None:
            raise ValueError('user_id is needed.')

        query = self._gen_query(user_id, from_date, to_date, filters)
        return query, {}, {'user_id': user_id, 'from': from_date, 'to': to_date, 'filters': filters}

    def _build_result(self, es_result, param):
        # keywords = param['keywords']
        # order = param['order']
        total = es_result['hits']['total']
        agg = es_result['aggregations']
        return {
            'user_id': param['user_id'],
            'from': param['from'],
            'to': param['to'],
            'post_count': total,
            "sum_likes": agg['sum_likes']['value'],
            "max_likes": agg['max_likes']['value'],
            "min_likes": agg['min_likes']['value'],
            "avg_likes": 0 if total == 0 else agg['sum_likes']['value'] / total,

            "sum_retweets": agg['sum_retweets']['value'],
            "max_retweets": agg['max_retweets']['value'],
            "min_retweets": agg['min_retweets']['value'],
            "avg_retweets": 0 if total == 0 else agg['sum_retweets']['value'] / total,

            "sum_comments": agg['sum_comments']['value'],
            "max_comments": agg['max_comments']['value'],
            "min_comments": agg['min_comments']['value'],
            "avg_comments": 0 if total == 0 else agg['sum_comments']['value'] / total,

            "sum_engagement": agg['sum_engagement']['value'],
            "max_engagement": agg['max_engagement']['value'],
            "min_engagement": agg['min_engagement']['value'],
            "avg_engagement": 0 if total == 0 else agg['sum_engagement']['value'] / total,
        }

    @staticmethod
    def _gen_filter(filters, key):
        filter_clause = []
        if key in filters:
            if filters[key]:
                values = filters[key]
                if isinstance(values, str):
                    values = values.split(' ')
                for fk in values:
                    filter_clause.append({'term': {key: fk}})
        return filter_clause

    @staticmethod
    def _gen_query(user_id, from_date, to_date, filters):
        query = {
            "query": {
                "bool": {
                    "filter":[
                        {'term': {'user_id': user_id}},
                        {
                            "range": {
                                        "publish_timestamp": {
                                            "from": from_date,
                                            "to": to_date
                                        }
                                    }
                        }
                    ]
                }
            },
            "aggs": {
                "sum_likes": {"sum": {"field": "likes", "missing": 0}},
                "max_likes": {"max": {"field": "likes", "missing": 0}},
                "min_likes": {"min": {"field": "likes", "missing": 0}},

                "sum_retweets": {"sum": {"field": "retweets", "missing": 0}},
                "max_retweets": {"max": {"field": "retweets", "missing": 0}},
                "min_retweets": {"min": {"field": "retweets", "missing": 0}},

                "sum_comments": {"sum": {"field": "comments", "missing": 0}},
                "max_comments": {"max": {"field": "comments", "missing": 0}},
                "min_comments": {"min": {"field": "comments", "missing": 0}},

                "sum_engagement": {"sum": {"field": "sum_engagement", "missing": 0}},
                "max_engagement": {"max": {"field": "sum_engagement", "missing": 0}},
                "min_engagement": {"min": {"field": "sum_engagement", "missing": 0}},
            },
            "size": 0
        }

        return query
